# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2024-07-15
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-07-15

# @Description: PDF Document Intelligence Loader for RAG

from typing import Iterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document
from langchain_core.messages import HumanMessage

from crewplus.services.docintel_service import DocintelService

from config.settings import *

class PDFDocintelLoader(BaseLoader):
    """An image document loader that loads an image file into Document."""

    def __init__(self, file_path: str, model: str = None) -> None:
        """Initialize the loader with a file path.

        Args:
            file_path: The path to the file to load.
        """
        self.file_path = file_path
        self.model = model

    def lazy_load(self) -> Iterator[Document]:  # <-- Does not take any arguments
        """A lazy loader that reads a file line by line.

        When you're implementing lazy load methods, you should use a generator
        to yield documents one by one.
        """
        
        docintel_srv = DocintelService(self.model)
    
        result = docintel_srv.analyze_document(self.file_path)

        yield Document(
            page_content=result.content,
            metadata={"file_type": "pdf", "source": self.file_path, "page": 0},
        )